import pickle as pickle
import pandas as pd
from pandas import DataFrame

import settings
import datetime
import helper

class Data:

    def __init__(self):
        self.raw_data = {}  #所有股票的数据
        self.selected_data = {}  #选中的股票的数据
        self.cleaned_data = {}     #数据清洗后的股票数据
        self.selected_stock_frame = None  # DataFrame

    def read_pickle_as_dict(self, data_path):
        """read given data as dict
        :get:
            dict{date:pandas.DataFrame}
        """
        with open(data_path, 'rb') as file:
            self.raw_data = pickle.load(file)

    def convert_as_csv(self, date_string):
        """
        将某日的原始数据输出为 csv 文件
        :param date_string: 
        :return:
        """
        self.read_pickle_as_dict(settings.DAILY_DATA_PATH)
        csv_path = '../data/' + date_string + '.csv'
        self.raw_data[date_string].to_csv(csv_path)

    def get_raw_data_between(self, start_date_string, end_date_string):
        """
        获取指定时间段的所有股票原始数据
        :return: dict{date_string: DataFrame}
        """
        self.selected_data = helper.select_data_between(start_date_string, end_date_string, self.raw_data)

    def prepare_data_for_factor(self, factor_list):
        """
        为单因子策略清洗数据
        :param factor_list: 因子所需索引列表
        :return: cleaned_dict
        """
        cleaned_dict = {}
        # index_list = self.selected_data[settings.DEFAULT_KEY].index.values
        
        #获取所有日期数据
        date_list = sorted(self.selected_data.keys())
        index_list = self.selected_data[date_list[0]].index.values
        # target_index_list = ['open', 'close', 'high', 'low']
        target_index_list = factor_list

        # for key, value in self.selected_data.items()
        for date in date_list:
            cleaned_data_frame = self.selected_data[date].dropna(axis=0, subset=target_index_list)
            index_list = list(set(index_list).intersection(set(cleaned_data_frame.index.values)))

        # 保留特定行和列
        for key, value in self.selected_data.items():
            raw_data_frame = value
            column_cleaned_data_frame = raw_data_frame[target_index_list]
            cleaned_data_frame = column_cleaned_data_frame[column_cleaned_data_frame.index.isin(index_list)]
            cleaned_dict[key] = cleaned_data_frame

        self.cleaned_data = cleaned_dict

    def get_stock_between(self,
                          start_date: str = '2010-01-05',
                          end_date: 'str' = '2011-01-05',
                          stock_number: 'str' = '600150.XSHG'):
        """
        获取指定时间段某只股票的数据
        :return: DataFrame{index: date_string, columns: factors}
        """
        new_data_frame: DataFrame = pd.DataFrame(columns=self.raw_data['2010-01-05'].loc['600150.XSHG'].index)
        date = start_date
        date_time = datetime.datetime.strptime(date, '%Y-%m-%d')
        end_date_time = datetime.datetime.strptime(end_date, '%Y-%m-%d')
        keys = self.raw_data.keys()
        while date_time <= end_date_time:
            if date in keys:
                new_data_frame.loc[date] = self.raw_data[date].loc[stock_number]
            date_time = date_time + datetime.timedelta(days=1)
            date = date_time.strftime("%Y-%m-%d")
        self.selected_stock_frame = new_data_frame

    def show_raw_data_format(self):
        """
        查看原始数据内容格式
        :return: 
        """
        print(f"读取后的数据格式为{type(self.raw_data)}")
        key, value = self.raw_data.popitem()
        print(f"字典中 key 的示例：{key}， 字典中 value 的类型：{type(value)}")
        df = self.raw_data["2010-01-05"]
        print("示例格式：")
        print(df.head(1))
        # print(df.index)
        # print(df.columns)
        # print(df.values)
